package com.gaga.common.utils;

/**
 * @DESCRIPTION:
 * @USER: ningxiang
 * @DATE: 4/26/2025 10:26 PM
 */

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.nio.file.StandardCopyOption;
import java.security.SecureRandom;
import java.util.List;
import java.util.UUID;

@Slf4j
public class FileCrypto {
    // 加密算法和参数
    private static final Path KEY_FILE = Paths.get("secret.key");
    private static final String ALGORITHM = "AES/CBC/PKCS5Padding";
    private static final String KEY_ALGORITHM = "AES";
    private static final int IV_LENGTH = 16; // AES块大小固定为16字节

    // 生成 AES 密钥
    public static void generateKey(Path keyPath) throws Exception {
        KeyGenerator keyGen = KeyGenerator.getInstance(KEY_ALGORITHM);
        keyGen.init(256);
        SecretKey key = keyGen.generateKey();
        Files.write(keyPath, key.getEncoded());
    }

    // 加密文件并记录原始文件名
    private static void encryptFile(Path sourceFile, Path targetDir, SecretKey key) throws Exception {
        String encryptedName = UUID.randomUUID() + ".enc";
        Path encryptedFile = targetDir.resolve(encryptedName);

        try (InputStream in = Files.newInputStream(sourceFile);
             DataOutputStream out = new DataOutputStream(Files.newOutputStream(encryptedFile))) {

            // 生成随机 IV
            byte[] iv = new byte[IV_LENGTH];
            new SecureRandom().nextBytes(iv);
            out.write(iv);

            // 写入原始文件名
            String originalName = sourceFile.getFileName().toString();
            byte[] nameBytes = originalName.getBytes("UTF-8");
            out.writeShort(nameBytes.length);
            out.write(nameBytes);

            // 加密内容
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.ENCRYPT_MODE, key, new IvParameterSpec(iv));
            try (CipherOutputStream cipherOut = new CipherOutputStream(out, cipher)) {
                byte[] buffer = new byte[4096];
                int bytesRead;
                while ((bytesRead = in.read(buffer)) != -1) {
                    cipherOut.write(buffer, 0, bytesRead);
                }
            }
        }

        // 删除源文件
        Files.delete(sourceFile);
    }

    // 解密文件并恢复原始文件名
    private static void decryptFile(Path encryptedFile, Path targetDir, SecretKey key) throws Exception {
        try (DataInputStream in = new DataInputStream(Files.newInputStream(encryptedFile))) {
            // 读取 IV
            byte[] iv = new byte[IV_LENGTH];
            in.readFully(iv);

            // 读取原始文件名
            int nameLength = in.readShort() & 0xFFFF;
            byte[] nameBytes = new byte[nameLength];
            in.readFully(nameBytes);
            String originalName = new String(nameBytes, "UTF-8");

            // 解密内容
            Path decryptedFile = targetDir.resolve(originalName);
            Cipher cipher = Cipher.getInstance(ALGORITHM);
            cipher.init(Cipher.DECRYPT_MODE, key, new IvParameterSpec(iv));
            try (CipherInputStream cipherIn = new CipherInputStream(in, cipher);
                 OutputStream out = Files.newOutputStream(decryptedFile)) {
                byte[] buffer = new byte[4096];
                int bytesRead;
                while ((bytesRead = cipherIn.read(buffer)) != -1) {
                    out.write(buffer, 0, bytesRead);
                }
            }
        }
    }

    // 新增modelName参数控制解密范围
    public static void processDirectory(Path sourceDir, Path targetDir, SecretKey key, boolean isEncrypt, List<String> modelNameList) throws Exception {
        Files.walk(sourceDir).forEach(sourcePath -> {
            try {
                Path relativePath = sourceDir.relativize(sourcePath);
                Path targetPath = targetDir.resolve(relativePath);

                if (Files.isDirectory(sourcePath)) {
                    Files.createDirectories(targetPath);
                    return;
                }

                String fileName = sourcePath.getFileName().toString();
                if (isEncrypt) {
                    // 加密逻辑保持不变
                    if (fileName.endsWith(".pth") || fileName.endsWith(".ckpt")) {
                        encryptFile(sourcePath, targetPath.getParent(), key);
                    } else {
                        Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
                    }
                } else {
                    String modelName = sourcePath.getParent().getFileName().toString();

                    if (CollectionUtils.isNotEmpty(modelNameList) && !modelNameList.contains(modelName)) {
                        return;
                    }
                    // 解密逻辑新增模型名称过滤
                    if (fileName.endsWith(".enc")) {
                        String baseName = fileName.substring(0, fileName.length() - 4);
                        Path decryptedPath = targetPath.resolveSibling(baseName);
                        if (!Files.exists(decryptedPath)) {
                            decryptFile(sourcePath, targetPath.getParent(), key);
                        } else {
                            log.info("跳过已存在的解密文件: " + decryptedPath);
                        }
                    } else {
                        Files.copy(sourcePath, targetPath, StandardCopyOption.REPLACE_EXISTING);
                    }
                }
            } catch (Exception e) {
                log.error("处理文件失败: " + sourcePath, e);
            }
        });
    }

    public static void main(String[] args) throws Exception {
        // 配置参数
        Path keyPath = Paths.get("secret.key");
        Path sourceDir = Paths.get("D:\\work\\video\\模型相关\\test");  // 源目录
        Path encryptedDir = Paths.get("D:\\work\\video\\模型相关\\test"); // 加密后目录
        Path decryptedDir = Paths.get("D:\\work\\video\\模型相关\\test\\fr"); // 解密后目录

        // 生成密钥（首次运行时执行）
        SecretKey key = getSecretKey();

        // 加密操作
//        processDirectory(sourceDir, encryptedDir, key, true);
//        System.out.println("加密完成，结果保存在: " + encryptedDir);

        // 解密操作
        processDirectory(encryptedDir, decryptedDir, key, false, null);
        System.out.println("解密完成，结果保存在: " + decryptedDir);
    }

    public static SecretKey getSecretKey() throws Exception {

        Path KEY_FILE = Paths.get("secret.key");
        log.info("KEY_FILE 1: " + KEY_FILE);
        Path externalKey = Paths.get(System.getProperty("user.dir"), "secret.key");
        log.info("KEY_FILE 2: " + externalKey);

        if (Files.exists(externalKey)) {
            KEY_FILE = externalKey;
        } else {
            // 尝试从 JAR 内资源加载
            try (InputStream is = FileCrypto.class.getResourceAsStream("/secret.key")) {
                if (is != null) {
                    Path tempFile = Files.createTempFile("secret", ".key");
                    Files.copy(is, tempFile, StandardCopyOption.REPLACE_EXISTING);
                    KEY_FILE = tempFile;
                    log.info("KEY_FILE 3: " + tempFile);
                }
            }
        }

        SecretKey key;
        if (!Files.exists(KEY_FILE)) {
            log.info("密钥文件不存在，正在生成密钥...");
            String password = "your-strong-password-here"; // 替换为实际密码
            key = AESKeyManager.generateFixedKey(password, KEY_FILE);
        } else {
            log.info("密钥文件存在，正在加载密钥...");
            key = AESKeyManager.loadFixedKey(KEY_FILE);
        }
        AESKeyManager.validateKey(key);
        return key;
    }
}